import torch
from config import *
from data_preprocessing import load_data, preprocess_advanced
from model import AdvancedLSTM
from model_trainer import ModelTrainer
from trading_simulator import TradingSimulator
from utils import check_data_distribution


def main_train():
    # 训练模式
    print("开始处理数据...")
    df = load_data(DATA_PATH)

    X_train, X_test, y_train, y_test, input_size = preprocess_advanced(
        df, lookback=LOOKBACK, forward_period=FORWARD_PERIOD, threshold=THRESHOLD
    )

    check_data_distribution(y_train, y_test)

    print("\n开始训练模型...")
    model = AdvancedLSTM(input_size=input_size, hidden_size=HIDDEN_SIZE,
                         num_layers=NUM_LAYERS, dropout=DROPOUT, num_classes=NUM_CLASSES)

    trainer = ModelTrainer(model)
    history = trainer.train(X_train, y_train, X_test, y_test,
                            num_epochs=NUM_EPOCHS, batch_size=BATCH_SIZE,
                            learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY,
                            patience=PATIENCE)

    # 最终评估
    test_preds, test_true = trainer.predict(X_test, y_test)
    trainer.generate_report(test_true, test_preds)

    print(f"随机猜测基准准确率: {1 / 3:.2%}")


def main_predict():
    # 预测模式
    print("开始处理数据...")
    df = load_data(DATA_PATH)

    X_train, X_test, y_train, y_test, input_size = preprocess_advanced(
        df, lookback=LOOKBACK, forward_period=FORWARD_PERIOD, threshold=THRESHOLD
    )

    _, test_data, _, _ = preprocess_advanced(
        df, lookback=LOOKBACK, forward_period=FORWARD_PERIOD,
        threshold=THRESHOLD, return_raw_data=True
    )

    print("加载已训练模型...")
    model = AdvancedLSTM(input_size=input_size, hidden_size=HIDDEN_SIZE,
                         num_layers=NUM_LAYERS, dropout=DROPOUT, num_classes=NUM_CLASSES)
    model.load_state_dict(torch.load(MODEL_PATH))

    trainer = ModelTrainer(model)
    predictions, _ = trainer.predict(X_test, y_test)

    # 交易模拟
    print("开始交易模拟...")
    simulator = TradingSimulator(
        initial_capital=INITIAL_CAPITAL,
        transaction_cost=TRANSACTION_COST,
        risk_per_trade=RISK_PER_TRADE,
        min_trade_interval=MIN_TRADE_INTERVAL
    )

    simulator.run_simulation(test_data, predictions, lookback=LOOKBACK)
    simulator.generate_performance_report()
    simulator.generate_simple_visualizations(price_data=test_data)


if __name__ == "__main__":
    # 根据需求选择训练或预测模式
    # main_train()  # 训练模型
    main_predict()  # 使用训练好的模型进行预测